import numpy as np
from numpy.random import random

class SVD_CF:
    def __init__(self, data, k=40):
        self.data = np.array(data)
        # 隐含变量的维度
        self.k = k
        #self.mu = np.mean(self.data[:, 2])  # average rating

        # item和用户的偏置项
        self.bi = {}
        self.bu = {}

        # item和用户的隐含向量
        self.qi = {}
        self.pu = {}

        self.ItemsForUser = {}  # 每个Item对应的用户
        self.UsersForItem = {}  # 每个用户对哪些Item打过分
        self.recommendForUser = {}
        self.recommendForTest = {}  # 测试数据集

        sum = 0
        items = 0
        for i in range(self.data.shape[0]):
            uid = self.data[i][0]  # user id
            i_id = self.data[i][1]  # item_id
            score = int(self.data[i][2])  # rating

            self.ItemsForUser.setdefault(i_id, {})
            self.UsersForItem.setdefault(uid, {})

            self.ItemsForUser[i_id][uid] = score
            self.UsersForItem[uid][i_id] = score

            self.bi.setdefault(i_id, 0)
            self.bu.setdefault(uid, 0)

            self.qi.setdefault(i_id, random((self.k, 1)) / 10 * (np.sqrt(self.k)))
            self.pu.setdefault(uid, random((self.k, 1)) / 10 * (np.sqrt(self.k)))

            sum += score
            items += 1

        self.mu = sum/items

    # 根据当前参数，预测用户uid对Item（i_id）的打分
    def predict(self, uid, i_id):
        self.bi.setdefault(i_id, 0)
        self.bu.setdefault(uid, 0)

        self.qi.setdefault(i_id, np.zeros((self.k, 1)))
        self.pu.setdefault(uid, np.zeros((self.k, 1)))

        if self.qi[i_id].all() is None:
            self.qi[i_id] = np.zeros((self.k, 1))
        if self.pu[uid].all() is None:
            self.pu[uid] = np.zeros((self.k, 1))

        ans = self.mu + self.bi[i_id] + self.bu[uid] + np.sum(self.qi[i_id] * self.pu[uid])
        return ans

    # gamma：为学习率
    # Lambda：正则参数
    def train(self, steps=50, gamma=0.1, Lambda=0.1):
        for step in range(steps):
            print('the ', step, '-th  step is running')
            rmse_sum = 0.0

            # 将训练样本打散顺序
            kk = np.random.permutation(self.data.shape[0])
            for j in range(self.data.shape[0]):
                # 每次一个训练样本
                i = kk[j]
                uid = self.data[i][0]
                i_id = self.data[i][1]
                rat = int(self.data[i][2])

                # 预测残差
                eui = rat - self.predict(uid, i_id)
                # 残差平方和
                rmse_sum += eui ** 2

                # 随机梯度下降，更新
                self.bu[uid] += gamma * (eui - Lambda * self.bu[uid])
                self.bi[i_id] += gamma * (eui - Lambda * self.bi[i_id])

                temp = self.qi[i_id]
                self.qi[i_id] += gamma * (eui * self.pu[uid] - Lambda * self.qi[i_id])
                self.pu[uid] += gamma * (eui * temp - Lambda * self.pu[uid])

                # 学习率递减
            gamma = gamma * 0.93
            print("the rmse of this step on train data is ", np.sqrt(rmse_sum / self.data.shape[0]))
            # self.test(test_data)

    def test(self, test_X):
        sums = 0
        sums2 = 0
        test_X = np.array(test_X)

        for i in range(test_X.shape[0]):  # 对每个测试样本
            user = test_X[i][0]
            item = test_X[i][1]
            score = int(test_X[i][2])

            self.recommendForTest.setdefault(user, {})
            self.recommendForUser.setdefault(user, {})
            self.recommendForTest[user][item] = score

            pre = self.predict(user, item)  # 预测打分
            self.recommendForUser[user][item] = pre

            sums += (pre - score) ** 2  # 残差平方和
            sums2 += abs(pre - score)
        rmse = np.sqrt(sums / test_X.shape[0])  # 均方误差
        mae = sums2 / test_X.shape[0]
        return rmse, mae

    def recommend(self, user, topN):
        list1 = sorted(self.recommendForUser[user].items(), key=lambda x: x[1], reverse=True)
        if topN > len(list1):
            return list1
        else:
            return list1[:topN]

    def precisionRecall(self, topN):
        hit = 0
        n_recall = 0
        n_precision = 0
        for user, items in self.recommendForUser.items():
            test = sorted(items.items(), key=lambda x: x[1], reverse=True)
            if topN > len(test):
                topN = len(test)
            testCollect = self.transformSet(test[:topN])
            recommendList = self.recommend(user, topN)
            recommendCollect = self.transformSet(recommendList)
            hit += len(testCollect.intersection(recommendCollect))
            n_recall += len(items)
            n_precision += topN
        return [hit / n_recall, hit / n_precision]

    def transformSet(self, data):
        collect = set()
        for i in data:
            collect.add(i[0])
        return collect
